import chromadb
import weaviate
from langchain_chroma import Chroma
from langchain_community.document_loaders import TextLoader
from langchain_text_splitters import CharacterTextSplitter
from langchain_weaviate import WeaviateVectorStore
from weaviate.auth import AuthApiKey

from langChain.config import model, embedding

print('-------------------- Chroma 向量数据库 ---------------------------------------')
# 加载文档并将其分割成片段
loader = TextLoader("../../resources/knowledge.txt", encoding="UTF-8")
documents = loader.load()
# 将其分割成片段
text_splitter = CharacterTextSplitter(chunk_size=1500, chunk_overlap=0)
docs = text_splitter.split_documents(documents)

# # 加载在 Chroma 内存中
# db = Chroma.from_documents(docs, embedding)
# # 进行查询
# docs = db.similarity_search("Pixar公司是做什么的?")
# # 打印结果
# print(docs[0].page_content)

# # 保存到磁盘
# db2 = Chroma.from_documents(docs, embedding, persist_directory="../../resource/vector_store/chroma_db")
# # 从磁盘加载
# db3 = Chroma(persist_directory="../../resource/vector_store/chroma_db", embedding_function=embedding)
# db3.as_retriever()
# docs = db3.similarity_search("Pixar公司是做什么的?")
# print(docs[0].page_content)


# # 创建一个chroma客户端
# persistent_client = chromadb.PersistentClient()
# # 创建一个集合（表）
# collection = persistent_client.get_or_create_collection("collection_1")
# collection.add(ids=["1", "2", "3"], documents=["a", "b", "c"])
# langchain_chroma = Chroma(
#     client=persistent_client,
#     collection_name="collection_1",
#     embedding_function=embedding,
# )
# print("在集合中有", langchain_chroma._collection.count(), "个文档")


# 创建简单的 ids
ids = [str(i) for i in range(1, len(docs) + 1)]
# 添加数据
example_db = Chroma.from_documents(docs, embedding, ids=ids)
docs = example_db.similarity_search("Pixar公司是做什么的?", k=1)
# #返回的距离分数是余弦距离。因此，分数越低越好。
# docs = example_db.similarity_search_with_score("Pixar公司是做什么的?")
# 更新文档的元数据
docs[0].metadata = {
    "source": "../../resource/knowledge.txt",
    "new_value": "hello world",
}

# print("更新前内容：", example_db._collection.get(ids=[ids[0]]))
# example_db.update_document(ids[0], docs[0])
# print("更新后内容：", example_db._collection.get(ids=[ids[0]]))
# # 删除最后一个文档
# print("删除前计数", example_db._collection.count())
# print(example_db._collection.get(ids=[ids[-1]]))
# example_db._collection.delete(ids=[ids[-1]])
# print("删除后计数", example_db._collection.count())
# print(example_db._collection.get(ids=[ids[-1]]))


# weaviate 向量数据库
# 连接到本地部署的 Weaviate
weaviate_client = weaviate.connect_to_custom(
    skip_init_checks=False,
    http_host="192.168.47.128",
    http_port=9001,
    http_secure=False,
    grpc_host="192.168.47.128",
    grpc_port=50051,
    grpc_secure=False,
    # 对应 AUTHENTICATION_APIKEY_ALLOWED_KEYS 中的密钥
    # 注意：此处只需要密钥即可，不需要用户名称
    auth_credentials=AuthApiKey("test-secret-key")
)

# 检查连接是否成功
print('是否连接成功 ：', weaviate_client.is_ready())

db = WeaviateVectorStore.from_documents(docs, embedding, client=weaviate_client)

docs = db.similarity_search_with_score("Pixar公司是做什么的?", k=1)
print('搜索到的文档 ： ', docs[0])

weaviate_client.close()
